# This file is part of Neuroinfo Toolkit.
#
# Neuroinfo Toolkit is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# Neuroinfo Toolkit is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with Neuroinfo Toolkit.  If not, see <http://www.gnu.org/licenses/>.

import subprocess
from neuro.base import Object
from neuro.exceptions import IllegalArgumentException
from neuro.exceptions import BaseException
import shlex
import time

def execute(command):
    '''
    Simple function to execute a command.
    
    :param command: Command
    :type command: str, list
    :return: dict {'status': int, 'stdout': str, 'stderr': str}
    '''
    if isinstance(command, basestring):
        com = shlex.split(command)
    else:
        com = command

    p = subprocess.Popen(com, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    (stdout, stderr) = p.communicate()

    out = {}
    out['status'] = p.returncode

    if(stdout.strip() == ''):
        out['stdout'] = None
    else:
        out['stdout'] = stdout.strip()
    
    if(stderr.strip() == ''):
        out['stderr'] = None
    else:
        out['stderr'] = stderr.strip()

    return out
    
class Command(Object):
    '''    
    This module provides a slightly more convenient way to execute a shell
    command with pipes, though it's not perfect ::
    
            >>> from neuro.command import Command
            >>> command = Command("echo 'foobar' | grep 'foobar'")
            >>> command.execute()
            >>> command.getStdout()
            'foobar'
    '''
    
    def __init__(self, command):
        '''
        Constructor ::
        
            >>> command = Command("echo 'foobar'")
            >>> command.execute()
            >>> command.getStdout()
            'foobar'
        '''
        self._command = None
        self._stdout = None
        self._stderr = None
        self._status = None
        self._pid = None
        self._process = None
        self._argc = 0
        
        self._setCommand(command)

    def _setCommand(self, command):
        '''
        Set command to be executed
        
        :param command: Command to be executed
        :type command: str, list
        '''
        ## --- input validation
        if(not isinstance(command, basestring) and not isinstance(command, list)):
            raise IllegalArgumentException("Command must be an instance of str or list")
        
        if(isinstance(command, basestring)):
            command = command.strip()
            
            if(command == ""):
                raise IllegalArgumentException("Command cannot be empty")
            
        self._command = command

    def toString(self):
        '''
        Get a shell-readable representation of the command :: 
        
            >>> command = Command(["ls", "-la", '/tmp'])
            >>> command.toString()
            'ls -la /tmp'
        
        :returns: String that can be cut-and-paste into a console
        :rtype: str
        '''
        if isinstance(self._command, basestring):
            return self._command

        cmd_str = ""

        for arg in self._command:
            temp_arg = arg
            if cmd_str != "":
                cmd_str += " "
            cmd_str += "'" + temp_arg.replace("'", "\\'") + "'"

        return cmd_str

    def getProvenanceXMLString(self):
        '''
        Get an XML representation of provenance information for this command. Not implemented.
        
        :returns: XML formatted provenance information
        :rtype: str
        '''
        return "<error>Not Implemented Yet</error>"

    def execute(self, fork=False, shell="/bin/bash", env=None):
        '''
        Execute command. Optionally, fork and detach process ::
        
            >>> command.execute()
        
        :param fork: Fork process
        :type fork: bool
        :param shell: Shell to execute the command with. If None, subprocess shell will be set to False
        :type shell: None, str
        :param env: Environment variables to pass to shell context
        :type env: dict
        :returns: Process ID
        :rtype: int
        :raises: :class:`CommandFailedException`
        '''
        executable = None

        if(isinstance(shell, basestring) and shell != ""):
    		if(shell.find("bash") > -1):
            		self._command = self._command + "; (exit ${PIPESTATUS})"
    
            	self._process = subprocess.Popen(self._command, shell=True, executable=shell, env=env,
                		stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    	else:
    		self._command = shlex.split(self._command)
    		self._argc = len(self._command)
    		self._process = subprocess.Popen(self._command, shell=False, env=env,
                            stdout=subprocess.PIPE, stderr=subprocess.PIPE)    
    
        if(not fork):
            self._start_time = time.time()
            (out, err) = self._process.communicate()

            if(out == None):
                self._stdout = ""
            else:
                self._stdout = out.strip()

            if(err == None):
                self._stderr = ""
            else:
                self._stderr = err.strip()

            self._status = self._process.returncode

            self._end_time = time.time()
            if(self._status > 0):
                raise CommandFailedException(self)

        self._pid = self._process.pid
        
    def getStatus(self):
        '''
        Retrieve command status code ::
        
            >>> status = command.getStatus()
            0
        
        :returns: Command status code
        :rtype: int
        '''
        return self._status
        
    def getCommand(self):
        '''
        Retrieve command ::
            
            >>> command = Command("ls -la")
            >>> command.getCommand()
            'ls -la'
        
        :returns: Command representation
        :rtype: str, list
        '''
        return self._command
        
    def getStdout(self):
        '''
        Retrieve executed command standard output ::
        
            >>> command = Command("echo 'foobar'")
            >>> command.execute()
            >>> command.getStdout()
            'foobar'
        
        :returns: Standard output
        :rtype: str
        '''
        return self._stdout
        
    def getStderr(self):
        '''
        Retrieve standard error ::
        
            >>> command = Command("echo 'foobar' 1>&2")
            >>> stderr = command.getStderr()
            'foobar'
        
        :returns: Standard error
        :rtype: str
        '''
        return self._stderr
        
    def getPid(self):
        '''
        Retrieve process ID
        
            >>> command.getPid()
            6867
        
        :returns: Process ID
        :rtype: int
        '''
        return self._pid

    def isAlive(self):
        '''
        If forked, check if the process is still alive

            >>> command.isAlive()
            True
        
        :rtype: bool
        '''
        if(self._process == None):
            return False

        if(self._process.poll() == None):
            return True
        else:
            self._status = self._process.returncode
            return False

class CommandFailedException(BaseException):
    '''
    Command execution exception
    '''

    def __init__(self, command):
        '''
        Constructor

        :param command: Command object
        :type command: :class:`Command`
        '''
        BaseException.__init__(self)

        if(not isinstance(command, Command)):
            raise IllegalArgumentException("Command must be an instance of Command")

        self._command = command

    def getCommand(self):
        '''
        Get the command that triggered this exception

        :rtype: :class:`Command`
        '''
        return self._command
    
    def getMessage(self):
        '''
        Get custom message
        
        :rtype: str
        '''
        message = "Command Failed\nCommand >> " + str(self._command.getCommand()) + "\n\n" + \
                "Exit Status >> " + str(self._command.getStatus()) +  "\n\n" + \
                "OUTPUT >> (see below)\n" + self._command.getStdout() + self._command.getStderr()

        return message
